/*=====================================================================*
 *                   Copyright (C) 2011 Paul Mineiro                   *
 * All rights reserved.                                                *
 *                                                                     *
 * Redistribution and use in source and binary forms, with             *
 * or without modification, are permitted provided that the            *
 * following conditions are met:                                       *
 *                                                                     *
 *     * Redistributions of source code must retain the                *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer.                                       *
 *                                                                     *
 *     * Redistributions in binary form must reproduce the             *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer in the documentation and/or            *
 *     other materials provided with the distribution.                 *
 *                                                                     *
 *     * Neither the name of Paul Mineiro nor the names                *
 *     of other contributors may be used to endorse or promote         *
 *     products derived from this software without specific            *
 *     prior written permission.                                       *
 *                                                                     *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND              *
 * CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,         *
 * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES               *
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE             *
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER               *
 * OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,                 *
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES            *
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE           *
 * GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR                *
 * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF          *
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT           *
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY              *
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE             *
 * POSSIBILITY OF SUCH DAMAGE.                                         *
 *                                                                     *
 * Contact: Paul Mineiro <paul@mineiro.com>                            *
 *=====================================================================*/

#include <boost/optional.hpp>
#include <boost/none.hpp>
#include "dyadicp.hh"

#include <cmath>

namespace
{
  class PFunc
    {
      private:
        const float adotb;
        const float anormplusbnorm;
        const float wdotx;
        const float xnorm;
        const float kappa;

        mutable bool larges_init;
        mutable float absadotb;
        mutable float absanormplusbnorm;
        mutable float abssum;
        mutable float signsum;
        mutable float logabssum;
        mutable float err_abssum;
        mutable float absdiff;
        mutable float signdiff;
        mutable float logabsdiff;
        mutable float err_absdiff;

        float p_smalls (float s) const;
        float p_mediums (float s) const;
        std::pair<bool, float> p_larges (float s) const;

        float pprime_smalls (float s) const;
        float pprime_mediums (float s) const;
        float pprime_larges (float s) const;

      public:
        PFunc (float _adotb,
               float _anormplusbnorm,
               float _wdotx,
               float _xnorm,
               float _kappa)
          : adotb (_adotb),
            anormplusbnorm (_anormplusbnorm),
            wdotx (_wdotx),
            xnorm (_xnorm),
            kappa (_kappa),
            larges_init (false),
            absadotb (0),
            absanormplusbnorm (0),
            abssum (0),
            signsum (0),
            logabssum (0),
            err_abssum (0),
            absdiff (0),
            signdiff (0),
            logabsdiff (0),
            err_absdiff (0)
          {
          }

        std::pair<bool, float>
        operator() (float s) const
          {
            std::pair<bool, float> rv;
  
            if (adotb == 0 && anormplusbnorm == 0)
              {
                rv.first = false;
                rv.second = wdotx - s * xnorm;
              }
            else 
              {
                float abss = fabsf (s);
              
                if (abss < 1e-1)
                  {
                    rv.first = false;
                    rv.second = p_smalls (s);
                  }
                else if (abss < 3)
                  {
                    rv.first = false;
                    rv.second = p_mediums (s);
                  }
                else
                  {
                    rv = p_larges (s);
                  }
              }
  
            return rv;
          }

        float
        fprime (float s) const
          {
            float rv;
  
            if (adotb == 0 && anormplusbnorm == 0)
              {
                rv = -xnorm;
              }
            else 
              {
                float abss = fabsf (s);
              
                if (abss < 1e-1)
                  {
                    rv = pprime_smalls (s);
                  }
                else if (abss < 3)
                  {
                    rv = pprime_mediums (s);
                  }
                else
                  {
                    rv = pprime_larges (s);
                  }
              }
  
            return rv;
          }
      };

  float
  PFunc::p_smalls (float s) const
    {
      float p = 0.0f;

      float coshpart[10];
      float kappapart[10];

      coshpart[0] = adotb;
      coshpart[1] = -anormplusbnorm;
      coshpart[2] = 2.0f * adotb;
      coshpart[3] = -2.0f / 3.0f * anormplusbnorm;
      coshpart[4] = 2.0f / 3.0f * adotb;

      if (s > 1e-2)
        {
          coshpart[5] = -2.0f / 15.0f * anormplusbnorm;
          coshpart[6] = 4.0f / 45.0f * adotb;
          coshpart[7] = -4.0f / 315.0f * anormplusbnorm;
          coshpart[8] = 2.0f / 315.0f * adotb;
          coshpart[9] = -2.0f / 2835.0f * anormplusbnorm;
        }

      float powkappa = kappa;

      kappapart[0] = 1.0f;
      kappapart[1] = 2.0f * powkappa;                   powkappa *= kappa;
      kappapart[2] = 2.0f * powkappa;                   powkappa *= kappa;
      kappapart[3] = 4.0f / 3.0f * powkappa;            powkappa *= kappa;
      kappapart[4] = 2.0f / 3.0f * powkappa;            powkappa *= kappa;

      if (s > 1e-2)
        {
          kappapart[5] = 4.0f / 15.0f * powkappa;       powkappa *= kappa;
          kappapart[6] = 4.0f / 45.0f * powkappa;       powkappa *= kappa;
          kappapart[7] = 8.0f / 315.0f * powkappa;      powkappa *= kappa;
          kappapart[8] = 2.0f / 315.0f * powkappa;      powkappa *= kappa;
          kappapart[9] = 4.0f / 2835.0f * powkappa;        
        }

      unsigned int ordermax = (s > 1e-2) ? 10 : 5;

      for (unsigned int pows = ordermax; pows > 0; --pows)
        {
          float coeff = 0.0f;

          for (unsigned int n = 0; n < pows; ++n)
            {
              coeff += coshpart[n] * kappapart[pows - n - 1];
            }

          p *= s;
          p += coeff;
        }

      p += wdotx - s * xnorm;

      return p;
    }

  float
  PFunc::pprime_smalls (float s) const
    {
      float p = 0.0f;

      float coshpart[10];
      float kappapart[10];

      coshpart[0] = adotb;
      coshpart[1] = -anormplusbnorm;
      coshpart[2] = 2.0f * adotb;
      coshpart[3] = -2.0f / 3.0f * anormplusbnorm;
      coshpart[4] = 2.0f / 3.0f * adotb;

      if (s > 1e-2)
        {
          coshpart[5] = -2.0f / 15.0f * anormplusbnorm;
          coshpart[6] = 4.0f / 45.0f * adotb;
          coshpart[7] = -4.0f / 315.0f * anormplusbnorm;
          coshpart[8] = 2.0f / 315.0f * adotb;
          coshpart[9] = -2.0f / 2835.0f * anormplusbnorm;
        }

      float powkappa = kappa;

      kappapart[0] = 1.0f;
      kappapart[1] = 2.0f * powkappa;           powkappa *= kappa;
      kappapart[2] = 2.0f * powkappa;           powkappa *= kappa;
      kappapart[3] = 4.0f / 3.0f * powkappa;    powkappa *= kappa;
      kappapart[4] = 2.0f / 3.0f * powkappa;    powkappa *= kappa;

      if (s > 1e-2)
        {
          kappapart[5] = 4.0f / 15.0f * powkappa;       powkappa *= kappa;
          kappapart[6] = 4.0f / 45.0f * powkappa;       powkappa *= kappa;
          kappapart[7] = 8.0f / 315.0f * powkappa;      powkappa *= kappa;
          kappapart[8] = 2.0f / 315.0f * powkappa;      powkappa *= kappa;
          kappapart[9] = 4.0f / 2835.0f * powkappa;        
        }

      unsigned int ordermax = (s > 1e-2) ? 10 : 5;

      for (unsigned int pows = ordermax; pows > 1; --pows)
        {
          float coeff = 0.0f;

          for (unsigned int n = 0; n < pows; ++n)
            {
              coeff += coshpart[n] * kappapart[pows - n - 1];
            }

          p *= (pows - 1) * s;
          p += coeff;
        }

      p -= xnorm;

      return p;
    }

  float
  PFunc::p_mediums (float s) const
    {
      // evaluate the straightforward way
      // p ==   1/2 * exp (2 * kappa * s) * adotb * ( exp (-2 s) + exp (2 s) )
      //      + 1/4 * exp (2 * kappa * s) * anormplusbnorm * ( exp (-2 s) - exp (2 s) )
      //      + wdotx - s * xnorm
      
      float exps = expf (2.0f * (kappa + 1.0f) * s);
      float expminuss = expf (2.0f * (kappa - 1.0f) * s);

      return 0.5f * adotb * (expminuss + exps)
           + 0.25f * anormplusbnorm * (expminuss - exps)
           + wdotx - s * xnorm;
    }

  float
  PFunc::pprime_mediums (float s) const
    {
      // evaluate the straightforward way
      
      float exps = expf (2.0f * (kappa + 1.0f) * s);
      float expminuss = expf (2.0f * (kappa - 1.0f) * s);

      return adotb * ((kappa - 1.0f) * expminuss + (kappa + 1.0f) * exps)
           + 0.5f * anormplusbnorm * 
               ((kappa - 1.0f) * expminuss - (kappa + 1.0f) * exps)
           - xnorm;
    }

  std::pair<bool, float>
  PFunc::p_larges (float s) const
    {
      // evaluate a convoluted way
      // p == 1/2 * (adotb + 1/2 * anormplusbnorm) * exp (-2 * (1 - kappa) * s)
      //      + 1/2 * (adotb - 1/2 * anormplusbnorm) * exp (2 * (1 + kappa) * s)
      //      + wdotx - s xnorm;
      //
      //   == exp (-2 * (1 - kappa) * s + log (abssum) - log (2)) signsum 
      //      + exp (2 * (1 - kappa) * s + log (absdiff) - log (2)) signdiff
      //      + wdotx - s xnorm;

      static const float logtwo = 0.69314718055994531f;
      std::pair<bool, float> rv;

      if (! larges_init)
        {
          absadotb = fabsf (adotb);
          absanormplusbnorm = fabsf (anormplusbnorm);

          float sum = adotb + 0.5f * anormplusbnorm;
          abssum = fabsf (sum);
          signsum = sum > 0 ? 1 : -1;
          if (abssum > 0)
            {
              logabssum = logf (abssum);
              err_abssum =
                1e-5f * (absadotb + 0.5f * absanormplusbnorm) / abssum;
            }
        
          float diff = adotb - 0.5f * anormplusbnorm;
          absdiff = fabsf (diff);
          signdiff = diff > 0 ? 1 : -1;
          if (absdiff > 0)
            {
              logabsdiff = logf (absdiff);
              err_absdiff = 
                1e-5f * (absadotb + 0.5f * absanormplusbnorm) / absdiff;
            }

          larges_init = true;
        }

      if (abssum > 0 && absdiff > 0)
        {
          rv.second = expf (-2.0f * (1.0f - kappa) * s + logabssum - logtwo) * signsum
                    + expf (2.0f * (1.0f + kappa) * s + logabsdiff - logtwo) * signdiff
                    + wdotx - s * xnorm;
          
          if (s > 0)
            {
              // we have to worry about diff precision
              
              // want:
              //   1/2 * err_absdiff * exp (2 * (1 + kappa) * s) << 1
              //   err_absdiff << 2 * exp (-2 * (1 + kappa) * s) 

              rv.first = (err_absdiff > 2.0 * expf (-2.0 * (1.0f + kappa) * s));
            }
          else if (s < 0)
            {
              // we have to worry about sum precision

              // want:
              //   1/2 * err_abssum * exp (-2 * (1 - kappa) * s) << 1
              //   err_abssum << 2 * exp (2 * (1 - kappa) * s) 
              //

              rv.first = (err_abssum > 2.0 * expf (2.0f * (1.0f - kappa) * s));
            }
        }
      else if (abssum > 0) // diff is zero to machine precision
        {
          rv.first = ((1.0f + kappa) * s > 0);
          rv.second = 
            expf (-2.0f * (1.0f - kappa) * s + logabssum - logtwo) * signsum
            + wdotx - s * xnorm;
        }
      else if (absdiff > 0) // sum is zero to machine precision
        {
          rv.first = ((1.0f - kappa) * s < 0);
          rv.second = 
            expf (2.0f * (1.0f + kappa) * s + logabsdiff - logtwo) * signdiff
            + wdotx - s * xnorm;
        }
      else // sum and diff are zero to machine precision
           // but adotb and anormplusbnorm are not identically zero
        {
          rv.first = true;
          rv.second = wdotx - s * xnorm;
        }

      return rv;
    }

  float
  PFunc::pprime_larges (float s) const
    {
      // evaluate a convoluted way
      // p == 1/2 * (adotb + 1/2 * anormplusbnorm) * exp (-2 * (1 - kappa) * s)
      //      + 1/2 * (adotb - 1/2 * anormplusbnorm) * exp (2 * (1 + kappa) * s)
      //      + wdotx - s xnorm
      //
      //   == exp (-2 s + log (abssum) - log (2)) signsum 
      //      + exp (2 s + log (absdiff) - log (2)) signdiff
      //      + wdotx - s xnorm

      static const float logtwo = 0.69314718055994531f;
      float rv;

      if (! larges_init)
        {
          absadotb = fabsf (adotb);
          absanormplusbnorm = fabsf (anormplusbnorm);

          float sum = adotb + 0.5f * anormplusbnorm;
          abssum = fabsf (sum);
          signsum = sum > 0 ? 1 : -1;
          if (abssum > 0)
            {
              logabssum = logf (abssum);
              err_abssum =
                1e-5f * (absadotb + 0.5f * absanormplusbnorm) / abssum;
            }
        
          float diff = adotb - 0.5f * anormplusbnorm;
          absdiff = fabsf (diff);
          signdiff = diff > 0 ? 1 : -1;
          if (absdiff > 0)
            {
              logabsdiff = logf (absdiff);
              err_absdiff = 
                1e-5f * (absadotb + 0.5f * absanormplusbnorm) / absdiff;
            }

          larges_init = true;
        }

      if (abssum > 0 && absdiff > 0)
        {
          rv = - 2.0f * expf (-2.0f * (1.0f - kappa) * s + logabssum - logtwo) * signsum
               + 2.0f * expf (2.0f * (1.0f + kappa) * s + logabsdiff - logtwo) * signdiff
               - xnorm;
        }
      else if (abssum > 0) // diff is zero to machine precision
        {
          rv = - 2.0f * expf (-2.0f * (1.0f - kappa) * s + logabssum - logtwo) * signsum
               - xnorm;
        }
      else if (absdiff > 0) // sum is zero to machine precision
        {
          rv = + 2.0f * expf (2.0f * (1.0f + kappa) * s + logabsdiff - logtwo) * signdiff
               - xnorm;
        }
      else // sum and diff are zero to machine precision
           // but adotb and anormplusbnorm are not identically zero
        {
          rv = -xnorm;
        }

      return rv;
    }
}

namespace flassol
{
namespace dyadic
{
  
std::pair<bool, float>
p (float s,
   float adotb,
   float anormplusbnorm,
   float wdotx,
   float xnorm,
   float kappa)
  {
    return PFunc (adotb, anormplusbnorm, wdotx, xnorm, kappa) (s);
  }

std::pair<bool, float>
pinv (float p,
      float adotb,
      float anormplusbnorm,
      float wdotx,
      float xnorm,
      float kappa)
  {
    PFunc f (adotb, anormplusbnorm, wdotx, xnorm, kappa);
    float s = 0;
    std::pair<bool, float> fs = std::make_pair (false, wdotx + adotb);
    float fsprime = -anormplusbnorm - xnorm + 2.0f * adotb * kappa;
    static const float mu = 0.00001f;

    for (unsigned int k = 0; 
         k < 10 && 
         fabsf (fs.second - p) > 9e-6 * (1.0 + fabs (fs.second) + fabs (p)); 
         ++k)
      {
        if (fabsf (fsprime) < 1e-8 * fabsf (fs.second - p)) // Newton fail
          {
            break;
          }

        float dir = (fs.second - p) / fsprime;
        float snew = s - dir;
        std::pair<bool, float> fsnew = f (snew);
        float fsnewprime = f.fprime (snew);

        for (unsigned int n = 0; 
             n < 14 &&
             (
                 // Armijo
                 (fsnew.second - p) * (fsnew.second - p) > 
                 (fs.second - p) * 
                  ( (fs.second - p) - 2.0f * mu * dir * fsprime )
             );
             ++n)
          {
            dir *= 0.25f;
            snew = s - dir;
            fsnew = f (snew);
            fsnewprime = f.fprime (snew);
          }

        if (fabsf (fsnew.second - p) > fabs (fs.second - p)) // line search fail
          {
            break;
          }

        s = snew;
        fs = fsnew;
        fsprime = fsnewprime;
      }

    std::pair<bool, float> rv;

    rv.first = 
      fabsf (fs.second - p) > 9e-6 * (1.0 + fabs (fs.second) + fabs (p)); 
    rv.second = s;

    return rv;
  }
}

}
